package com.gitee.wsl.struct.tree.ext

import com.gitee.wsl.struct.tree.sample.link.LinkTreeElement
import com.gitee.wsl.struct.tree.NodeElement
import com.gitee.wsl.struct.tree.Tree
import com.gitee.wsl.struct.tree.TreeElement

private fun<T,E: TreeElement<T, E>> Tree<T, E>.walk(breathFirst: Boolean) = sequence<E>{
    val queue = children.toMutableList()
    while (queue.isNotEmpty()) {
        val next = queue.removeFirst()
        yield(next)
        if (next is NodeElement<*, *>) {
            if(next is LinkTreeElement.Node<*>)
                next.open()
            next as NodeElement<T, E>
            if (breathFirst) {
                queue.addAll(next.children.orEmpty())
            } else {
                queue.addAll(0, next.children.orEmpty())
            }
        }
    }
}

fun<T,E: TreeElement<T, E>> Tree<T, E>.walkBreadthFirst(): Sequence<E> = walk(true)

fun<T,E: TreeElement<T, E>> Tree<T, E>.walkDepthFirst(): Sequence<E> = walk(false)